# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import argparse
import re
import json
import ast
import uvloop

from vllm.entrypoints.openai.api_server import run_server
from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args

from openmind.utils import logging
from openmind.flow.model.loader import get_init_kwargs


logger = logging.get_logger()
logging.set_verbosity_info()


class DeployvLLM:
    def __init__(self, args: argparse.Namespace):
        if args.backend_config:
            vllm_args = self._parse_args(args.backend_config)
        else:
            vllm_args = {}

        _ = get_init_kwargs()
        vllm_args["model"] = args.model_name_or_path
        vllm_args["host"] = getattr(vllm_args, "host", "127.0.0.1")
        vllm_args["port"] = getattr(args, "port", getattr(vllm_args, "port", 1025))

        parser = argparse.ArgumentParser()
        parser = make_arg_parser(parser)
        self.args = parser.parse_args(args=[], namespace=argparse.Namespace(**vllm_args))

        validate_parsed_serve_args(self.args)

    @staticmethod
    def _parse_args(backend_config_str: str) -> dict:
        """Parse the backend_config string into a dictionary.
        Args:
            backend_config_str (str): A string containing comma-separated key-value pairs.
        Returns:
            dict: A dictionary of parsed parameters with reasonable data types.
        """
        pattern = re.compile(r"([\w-]+)=({.*?}|\[.*?\]|[^,]+)")
        matches = pattern.findall(backend_config_str)

        parsed_config = {}
        for key, value in matches:
            key = key.strip().replace("-", "_")
            value = value.strip()

            if value.startswith("{") or value.startswith("["):
                try:
                    parsed_value = json.loads(value)
                    parsed_config[key] = parsed_value
                    continue
                except json.JSONDecodeError:
                    pass

            try:
                parsed_value = ast.literal_eval(value)
            except (ValueError, SyntaxError):
                parsed_value = value

            parsed_config[key] = parsed_value

        return parsed_config

    def deploy(self):
        uvloop.run(run_server(self.args))
